import os
import numpy as np
from osgeo import gdal
import cv2
import math
from multiprocessing import Pool
import time


class GRID:
    def load_image(self, filename):
        image = gdal.Open(filename)
 
        img_width = image.RasterXSize
        img_height = image.RasterYSize
 
        img_geotrans = image.GetGeoTransform()
        img_proj = image.GetProjection()
        img_data = image.ReadAsArray(0, 0, img_width, img_height)
        del image
 
        return img_proj, img_geotrans, img_data
 
    def write_image(self, filename, img_proj, img_geotrans, img_data):
        # 判断栅格数据类型
        if 'int8' in img_data.dtype.name:
            datatype = gdal.GDT_Byte
        elif 'int16' in img_data.dtype.name:
            datatype = gdal.GDT_UInt16
        else:
            datatype = gdal.GDT_Float32
 
        # 判断数组维度
        if len(img_data.shape) == 3:
            img_bands, img_height, img_width = img_data.shape
        else:
            img_bands, (img_height, img_width) = 1, img_data.shape
 
        # 创建文件
        driver = gdal.GetDriverByName('GTiff')
        image = driver.Create(filename, img_width, img_height, img_bands, datatype)
 
        image.SetGeoTransform(img_geotrans)
        image.SetProjection(img_proj)
 
        if img_bands == 1:
            image.GetRasterBand(1).WriteArray(img_data)
        else:
            for i in range(img_bands):
                image.GetRasterBand(i+1).WriteArray(img_data[i])
 
        del image # 删除变量,保留数据


class Graph:
    def __init__(self, file_path):
        self.file_path = os.path.join(file_path, 'LOD')
        layers = os.listdir(self.file_path)
        # 剔除非数字的文件夹
        layers = [layer for layer in layers if layer.isdigit()]
        sorted_layers = sorted(layers, key=lambda x: int(x))
        self.current_img = None
        self.layers = self.generate_layer_data(sorted_layers)
        self.height, self.width = cv2.imread(self.layers[-1][0][0]).shape[:2]
        self.current_layer = len(self.layers) - 1
        
    def generate_layer_data(self, sorted_layers:list):
        layers = []
        for layer in sorted_layers:
            layer_path = os.path.join(self.file_path, layer)
            images = os.listdir(layer_path)
            images = [image for image in images if image.endswith('.jpg')]
            i_num = 0
            j_num = 0
            for image in images:
                j, i = image.split('.')[0].split('_')
                i_num = max(i_num, int(i))
                j_num = max(j_num, int(j))
            # 如果是最后一层
            if i_num == 0 or j_num == 0:
                self.current_img = [(0, 0), (i_num+1, j_num+1)]
            # 构建一个二维数组
            layer_data = []
            for i in range(i_num+1):
                row = []
                for j in range(j_num+1):
                    row.append(os.path.join(layer_path, f'{j}_{i}.jpg'))
                layer_data.append(row)
            layers.append(layer_data)
        return layers
    
    def get_block_num(self):
        return (self.current_img[1][0] - self.current_img[0][0]) * (self.current_img[1][1] - self.current_img[0][1])
    
    def get_origin_size(self):
        width = self.width * len(self.layers[0][0])
        height = self.height * len(self.layers[0])
        return width, height
    
def get_LOD_part(sub_image, i, j, output_layer, block_sizeX, block_sizeY, contain_pixelX, contain_pixelY):
    sub_image = sub_image.transpose(2,1,0)
    # 由RGB转BGR    
    sub_image = cv2.cvtColor(sub_image, cv2.COLOR_RGB2BGR)
    sub_image = cv2.resize(sub_image, (block_sizeX, block_sizeY))
    cv2.imwrite(os.path.join(output_layer, f'{i}_{j}.jpg'), sub_image)
    print(f'Processing {i}_{j}.jpg')

def get_LOD_part2(input_layer, i, j, output_layer, block_sizeX, block_sizeY):
    sub_image1 = cv2.imread(os.path.join(input_layer, f'{2*i}_{2*j}.jpg'))
    sub_image2 = cv2.imread(os.path.join(input_layer, f'{2*i+1}_{2*j}.jpg'))
    sub_image3 = cv2.imread(os.path.join(input_layer, f'{2*i}_{2*j+1}.jpg'))
    sub_image4 = cv2.imread(os.path.join(input_layer, f'{2*i+1}_{2*j+1}.jpg'))
    # 拼接图片
    sub_image = np.vstack((np.hstack((sub_image1, sub_image2)), np.hstack((sub_image3, sub_image4))))
    sub_image = cv2.resize(sub_image, (block_sizeX, block_sizeY))
    cv2.imwrite(os.path.join(output_layer, f'{i}_{j}.jpg'), sub_image)
    # print(f'Processing {i}_{j}.jpg')

def error_callback(e):
    print(e)


def image_LOD(input_file, output_path,patch_size_x=32, patch_size_y=32, process_bar=None):
    '''
    input_file: 输入文件路径
    output_path: 输出文件路径
    patch_size: 将原图的每条边分割成多少块，必须为2的整数次幂
    '''
    if process_bar:
        process_bar.emit("读取图片", 0)
    start = time.time()
    file_name = input_file.split('/')[-1].split('.')[:-1]
    file_name = '.'.join(file_name)
    proj, geotrans, data = GRID().load_image(input_file)
    # 转置最后两维
    data = np.transpose(data,axes=(0,2,1))
    channel, width, height = data.shape
    output_dir = os.path.join(output_path, 'LOD')
    k = 0
    block_sizeX, block_sizeY = width // patch_size_x, height // patch_size_y
    pool = Pool(processes=16)
    num = math.log2(min(patch_size_x, patch_size_y)) + 1
    interval = int(80 // num)
    # print(interval)
    current = 20
    while 2 ** k <= min(patch_size_x, patch_size_y):
        if process_bar:
            process_bar.emit(f"处理第{k}级", current)
            current += interval
        output_layer = os.path.join(output_dir, str(k))
        if not os.path.exists(output_layer):
            os.makedirs(output_layer)
        contain_pixelX = width // (patch_size_x // 2 ** k)
        contain_pixelY = height // (patch_size_y // 2 ** k)
        if k == 0:
            # res = []
            for i in range(width//contain_pixelX):
                for j in range(height//contain_pixelY):
                    print(f'Processing {i}_{j}.jpg')
                    sub_image = data[:, i*contain_pixelX:(i+1)*contain_pixelX, j*contain_pixelY:(j+1)*contain_pixelY]
                    # res.append(pool.apply_async(get_LOD_part, args=(sub_image, i, j, output_layer, block_sizeX, block_sizeY, contain_pixelX, contain_pixelY), error_callback=error_callback))
                    sub_image = sub_image.transpose(2,1,0)
                    print(f"sub_image shape: {sub_image.shape}")
                    # 由RGB转BGR    
                    sub_image = cv2.cvtColor(sub_image, cv2.COLOR_RGB2BGR)
                    sub_image = cv2.resize(sub_image, (block_sizeX, block_sizeY))
                    cv2.imwrite(os.path.join(output_layer, f'{i}_{j}.jpg'), sub_image)
                    print(f'output_layer: {output_layer}')
            # for r in res:
            #     r.wait()
        else:
            res = []
            input_layer = os.path.join(output_dir, str(k-1))
            for i in range(width//contain_pixelX):
                for j in range(height//contain_pixelY):
                    res.append(pool.apply_async(get_LOD_part2, (input_layer, i, j, output_layer, block_sizeX, block_sizeY)))
                    # sub_image1 = cv2.imread(os.path.join(input_layer, f'{2*i}_{2*j}.jpg'))
                    # sub_image2 = cv2.imread(os.path.join(input_layer, f'{2*i+1}_{2*j}.jpg'))
                    # sub_image3 = cv2.imread(os.path.join(input_layer, f'{2*i}_{2*j+1}.jpg'))
                    # sub_image4 = cv2.imread(os.path.join(input_layer, f'{2*i+1}_{2*j+1}.jpg'))
                    # # 拼接图片
                    # sub_image = np.vstack((np.hstack((sub_image1, sub_image2)), np.hstack((sub_image3, sub_image4))))
                    # sub_image = cv2.resize(sub_image, (block_sizeX, block_sizeY))
                    # cv2.imwrite(os.path.join(output_layer, f'{i}_{j}.jpg'), sub_image)
            for r in res:
                r.wait()
        k += 1
    pool.close()
    pool.join()
    if process_bar:
        process_bar.emit("完成", 100)
    end = time.time()
    print(f'Finished in {end-start} seconds')
            
